from tensorflow import keras
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.layers import Dense,Softmax,ReLU,Conv2D,Flatten,MaxPool2D
import time
import matplotlib.pyplot as plt
import os
class Network(object):
    def __init__(self,train_x,train_y,test_x,test_y):
        self.train_x = train_x
        self.train_y = train_y
        self.test_x = test_x
        self.test_y = test_y
        self.save_model = 'train_cnn.h5'
        self.model = self.build()
        self.labels = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
                       'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
    def build(self):
        if os.path.exists(self.save_model):
            model = keras.models.load_model(self.save_model)
            return model
        else:
            model = Sequential()
            model.add(Conv2D(32, (3, 3), padding='same', activation='relu',
                      name = 'conv1', input_shape=(28, 28, 1)))
            model.add(MaxPool2D(pool_size=(2,2),name = 'pool1'))


            model.add(Conv2D(64, 3, strides=(1, 1), padding='valid',activation='relu',name='conv2'))
            model.add(MaxPool2D(pool_size=(2,2), strides=(2, 2), padding='same',name='pool2'))


            model.add(Conv2D(128, 3, strides=(1, 1), padding='valid', activation='relu',name='conv3'))
            model.add(MaxPool2D(pool_size=(2, 2), strides=(2, 2), padding='valid',name='pool3'))


            model.add(Flatten(name='fc'))
            model.add(Dense(10,activation='softmax',name='output'))

            opt = RMSprop(lr=0.0001,decay=0.000001)
            model.compile(loss='sparse_categorical_crossentropy',optimizer=opt,metrics=['accuracy'])
            model.summary()
            model.fit(self.train_x.reshape(-1,28,28,1), self.train_y, batch_size=256, epochs=60,validation_data=(self.test_x.reshape(-1,28,28,1),self.test_y))
            model.save('train_cnn.h5')
            self.model = model
            return model
    def predict(self,random_sample = 16):
        rand_data = np.random.choice(1000, random_sample)
        count = 0
        test_image = []
        test_label = []
        for i in rand_data:
            temp_x = self.test_x[i,:,:]
            test_image.append(temp_x)

            temp_y = self.test_y[i]
            test_label.append(temp_y)
        test_image = np.array(test_image)
        start_time = time.clock()
        result_predict = self.model.predict_classes(test_image.reshape(-1,28,28,1)).tolist()
        end_time = time.clock()
        print("Predict Image nums:{},elapse time:{:.4f}".format(random_sample,end_time-start_time))
        # result_predict = self.model.predict_classes(test_image).tolist()
        fig_col = int(np.sqrt(random_sample))
        fig_row = random_sample // fig_col
        for i in range(random_sample):
            if test_label[i] == result_predict[i]:
                count += 1
            predict_label = self.labels[result_predict[i]]
            true_label = self.labels[test_label[i]]
            plt.subplot(fig_row,fig_col,i+1)
            plt.imshow(test_image[i],cmap=plt.cm.binary)
            plt.title(predict_label,fontsize=4,color='red')
            plt.xticks(fontsize=4)
            plt.yticks(fontsize=4)
            # plt.box(False)
            plt.xticks([])
            plt.yticks([])
            # print('Predict Label:{},True label:{}'.format(
            #     predict_label,true_label))
        # print('Acuracy:{:.4f},Error nums:{}'.format(count/random_sample,random_sample-count))
        plt.savefig('../../figures/mnist_fashion_cnn_predict.png', dpi=600)
def main():
    fashion_mnist = keras.datasets.fashion_mnist
    (train_x,train_y),(test_x,test_y) = fashion_mnist.load_data()
    net = Network(train_x,train_y,test_x,test_y)
    net.predict()
if __name__ == '__main__':
    main()

